package com.ztools.common.util.common.aspect;

import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import com.ztools.common.util.common.Result;
import com.ztools.common.util.common.annotateUnit.AccessLimit;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import java.util.Map;

/**
 * @Author: seeTools
 * @Date: 2023/6/13  8:57
 * @Version 1.0
 * 限流AOP
 */
@Scope
@Aspect
@Component
@Slf4j
public class RateLimitAspect {

    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Pointcut("@annotation(com.ztools.common.util.common.annotateUnit.AccessLimit)")
    private void pointcut() {
    }

    @Around(value = "pointcut()")
    public Object around(ProceedingJoinPoint point) {
        Object result = null;
        try {
            // 获取注解
            AccessLimit accessLimit = ((MethodSignature) point.getSignature()).getMethod().getAnnotation(AccessLimit.class);
            // 注解不为空
            if (accessLimit != null) {
                // 获取限流令牌
                String limitKey = accessLimit.key();
                RateLimiter rateLimiter;
                if (!limitMap.containsKey(limitKey)) {
                    // 创建令牌桶
                    rateLimiter = RateLimiter.create(accessLimit.permitsPerSecond());
                    limitMap.put(limitKey, rateLimiter);
                    log.info("新建了令牌桶={},容量={}", limitKey, accessLimit.permitsPerSecond());
                } else {
                    rateLimiter = limitMap.get(limitKey);
                }
                // 获取令牌
                boolean acquire = rateLimiter.tryAcquire(accessLimit.timeout(), accessLimit.timeUnit());
                if (!acquire) {
                    log.warn("获取令牌失败={}", limitKey);
                    result = Result.error(accessLimit.msg());
                } else {
                    log.info("获取令牌成功={}", limitKey);
                    result = point.proceed();
                }
            }
        } catch (Throwable e) {
            e.printStackTrace();
        }
        return result;
    }
}
